import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

class Ploter:
    def __init__(self, device_file_path: str, agent_file_path: str):
        # Read data
        self.device_file_path = device_file_path
        self.agent_file_path  = agent_file_path
        self.device_data = pd.read_csv(device_file_path, header=None)
        self.agent_data  = pd.read_csv(agent_file_path, header=None)
        # Create subplots
        self.fig = plt.figure()
        self.ax = []
        for i in range(4):
            self.ax.append(self.fig.add_subplot(2, 2, i+1))
        plt.tight_layout()
        plt.show(block=False)
    
    def update(self):
        # Read data
        self.device_data = pd.read_csv(self.device_file_path, header=None)
        self.agent_data  = pd.read_csv(self.agent_file_path, header=None)
        # Update figures
        self.plot_freq()
        self.plot_temp()
        self.plot_fps()
        self.plot_loss()
        self.draw()

    def plot_freq(self):
        self.ax[0].clear()
        self.plot(self.ax[0], self.device_data, 0, 'CPU_l')
        self.plot(self.ax[0], self.device_data, 4, 'CPU_m')
        self.plot(self.ax[0], self.device_data, 7, 'CPU_b')
        self.add_info(self.ax[0], title='Frequencies', x_label='time', y_label='frequencies')

    def plot_temp(self):
        self.ax[1].clear()
        self.plot(self.ax[1], self.device_data, 9, 'Temp')
        self.add_info(self.ax[1], title='Temperature', x_label='time', y_label='temperature')
    
    def plot_fps(self):
        self.ax[2].clear()
        self.plot(self.ax[2], self.device_data, 10, 'FPS')
        self.add_info(self.ax[2], title='FPS', x_label='time', y_label='FPS')
    
    def plot_loss(self):
        self.ax[3].clear()
        self.plot(self.ax[3], self.agent_data, 1, 'loss')
        self.add_info(self.ax[3], title='Training Loss', x_label='time', y_label='loss')

    def plot(self, ax, data, row_index: int, label: str):
        y = data.iloc[row_index]
        x = np.arange(1, len(y)+1)
        ax.plot(x, y, label=label)
        xticks = np.arange(1, len(y)+1, 20)
        # ax.xticks(xticks)
        # ax.legend()

    def add_info(self, ax, title: str, x_label: str, y_label: str):
        ax.set_title(title)
        # plt.title(title)
        # ax.xlabel(x_label)
        # ax.ylabel(y_label)

    def draw(self):
        plt.legend()
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
        # plt.show()

class AgentPloter:
    def __init__(self, row, col, num, title):
        # Create subplots
        self.fig = plt.figure()
        self.ax = []
        self.title = title
        for i in range(num):
            self.ax.append(self.fig.add_subplot(row, col, i+1))
        plt.tight_layout()
        plt.show(block=False)
    
    def plot(self, ax, data):
        x = np.arange(1, len(data)+1)
        ax.plot(x, data)

    def draw(self):
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()

    def update(self, *args):
        '''
        Update figures
        '''
        id = 0
        for data in args:
            self.ax[id].clear() # clear
            self.ax[id].set_title(self.title[id])
            self.plot(self.ax[id], data) # redraw
            id += 1
        self.draw()





if __name__ == '__main__':

    ploter = Ploter('231031-log_device_test_200.csv', '231031-log_agent_200.csv')

    # 持续更新子图
    for i in range(1):
        ploter.update()

        
    os.system('pause')
    # 关闭图形窗口
    # plt.close()